import os
import cv2
from distutils.dir_util import copy_tree


class GenerateData:
    """
    生成训练集、验证集、测试集

    配置参数：配置所需的参数放在该类中的全局变量处
    """
    path_mt_human = '../../data/matting_human'  # matting_human数据集路径
    path_save = '../../data/datasets'  # 数据集保存路径
    path_debug = '../../data/datasets_debug'  # 调试用数据集保存路径
    proportion = '6:2:2'  # 训练集、验证集、测试集划分比例

    # 操作步骤
    COPY_PIC = False  # 是否创建文件夹并拷贝图片数据
    GENERATE_MASK = False  # 是否生成matting图对应的mask图
    GENERATE_TXT = True  # 是否生成matting图对应的TXT文件
    GENERATE_TRIMAP = False  # 是否生成由mask图生成trimap图
    KSIZE = 10  # 生成trimap图时指定的卷积核大小
    GENERATE_ALPHA = False  # 是否复制mask图为alpha
    GENERATE_DEBUG = True  # 是否生成调试用数据集

    # 已发现错误文件，需要删除的列表
    list_error_files = ['{}/matting/1803201916/._matting_00000000'.format(path_mt_human),
                        '{}/clip_img/1803241125/clip_00000000/._1803241125-00000005.jpg'.format(path_mt_human)]

    def run(self):
        """
        生成数据集的入口
        """
        self.remove_error_files()
        self.split_datasets_and_save()
        self.generate_mask()
        self.generate_txt(self.path_save)
        self.generate_trimap()
        self.generate_alpha()
        self.generate_debug()

    def remove_error_files(self):
        """
        删除错误文件
        """
        for error_file in self.list_error_files:
            if os.path.exists(error_file):
                os.remove(error_file)

    @staticmethod
    def safe_makedirs(path_dir):
        if not os.path.exists(path_dir):
            os.makedirs(path_dir)

    def split_datasets_and_save(self):
        """
        拆分并保存数据集
        """
        # 统计数据集数量
        sums = 0
        dict_time_count = dict()
        list_time_stamp = os.listdir('{}/matting'.format(self.path_mt_human))
        for time_stamp in list_time_stamp:

            path_time_stamp = '{}/matting/{}'.format(self.path_mt_human, time_stamp)
            list_matting = os.listdir(path_time_stamp)

            print(time_stamp)
            dict_time_count[time_stamp] = 0
            for pic in list_matting:
                path_pic = os.path.join(path_time_stamp, pic)

                if not os.path.isdir(path_pic):  # 路径类型检测, 错误：['../matting_human/matting/1803201916/._matting_00000000']
                    continue
                # for pic_ in os.listdir(path_pic):  # 图片名称检测, 错误:[._1803241125-00000005.jpg]
                #     if pic_.split('.')[1] not in ['jpg', 'png']:
                #         print('error images is []'.format(pic_))
                #         path_img = os.path.join(path_pic, pic_)
                #         os.remove(path_img)

                nums = len(os.listdir(path_pic))
                dict_time_count[time_stamp] += nums
                sums += nums

                print('——', pic, nums)

        print(dict_time_count)
        dt = sorted(dict_time_count.items(), key=lambda it: it[1], reverse=True)
        print(dt)

        # 训练集、验证集、测试集
        s = 0
        p = list()
        for x in self.proportion.split(':'):
            s += sums * int(x) // 10
            p.append(s)
        print(p)

        # 记录划分的数据集
        list_train = list()
        list_eval = list()
        list_test = list()
        sum_count = 0
        for item in dt:
            sum_count += item[1]
            if sum_count < p[0]:
                list_train.append(item[0])
            elif sum_count > p[1]:
                list_test.append(item[0])
            else:
                list_eval.append(item[0])

        # 创建文件夹并拷贝图片数据
        dir_names = ['train', 'eval', 'test']
        if self.COPY_PIC:
            for dirs in dir_names:

                path_mattiing = '{}/{}/matting'.format(self.path_save, dirs)
                self.safe_makedirs(path_mattiing)
                for ts in locals()['list_' + dirs]:
                    path_src = '{}/matting/{}'.format(self.path_mt_human, ts)
                    path_dst = '{}/{}'.format(path_mattiing, ts)
                    copy_tree(path_src, path_dst)

                path_clip_img = '{}/{}/clip_img'.format(self.path_save, dirs)
                self.safe_makedirs(path_clip_img)
                for ts in locals()['list_' + dirs]:
                    path_src = '{}/clip_img/{}'.format(self.path_mt_human, ts)
                    path_dst = '{}/{}'.format(path_clip_img, ts)
                    copy_tree(path_src, path_dst)

    def generate_mask(self):
        """
        生成mask文件夹
        """
        dir_names = ['train', 'eval', 'test']
        if self.GENERATE_MASK:
            for dirs in dir_names:
                path_mattiing = '{}/{}/matting'.format(self.path_save, dirs)
                path_mask = '{}/{}/mask'.format(self.path_save, dirs)
                self.safe_makedirs(path_mask)

                for time_stamp in os.listdir(path_mattiing):
                    path_ts = os.path.join(path_mattiing, time_stamp)
                    for mt in os.listdir(path_ts):
                        path_mt = os.path.join(path_ts, mt)

                        path_mask_mt = '{}/{}/mask/{}/{}'.format(self.path_save, dirs, time_stamp, mt)
                        self.safe_makedirs(path_mask_mt)

                        for pic in os.listdir(path_mt):
                            path_pic_src = os.path.join(path_mt, pic)
                            path_pic_dst = os.path.join(path_mask_mt, pic)

                            in_image = cv2.imread(path_pic_src, cv2.IMREAD_UNCHANGED)
                            alpha = in_image[:, :, 3]
                            cv2.imwrite(path_pic_dst, alpha)

    def generate_txt(self, path_save):
        """
        生成txt文件

        Ages:
            path_save: txt文件以 path_save/[train, eval, test]/[train, eval, test].txt 的形式保存
        """
        dir_names = ['train', 'eval', 'test']
        if self.GENERATE_TXT or self.GENERATE_DEBUG:
            for dirs in dir_names:
                file_path = '{}/{}/{}.txt'.format(path_save, dirs, dirs)
                with open(file_path, 'w', encoding='UTF-8') as f:
                    path_mattiing = '{}/{}/matting'.format(path_save, dirs)
                    for time_stamp in os.listdir(path_mattiing):
                        path_ts = os.path.join(path_mattiing, time_stamp)
                        for mt in os.listdir(path_ts):
                            path_mt = os.path.join(path_ts, mt)
                            for pic in os.listdir(path_mt):
                                path_pic = '{}/{}/{}'.format(time_stamp, mt, pic)
                                f.write(path_pic + "\n")
                print('{}.txt is generated！'.format(dirs))

    @staticmethod
    def erode_dilate(msk, struc="ELLIPSE", size=(10, 10)):
        """
        腐蚀膨胀操作
        """
        if struc == "RECT":
            kernel = cv2.getStructuringElement(cv2.MORPH_RECT, size)
        elif struc == "CORSS":
            kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, size)
        else:
            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, size)

        msk = msk / 255

        dilated = cv2.dilate(msk, kernel, iterations=1) * 255
        eroded = cv2.erode(msk, kernel, iterations=1) * 255

        res = dilated.copy()
        res[((dilated == 255) & (eroded == 0))] = 128

        return res

    def generate_trimap(self):
        """
        由mask图生成trimap图
        """
        dir_names = ['train', 'eval', 'test']
        if self.GENERATE_TRIMAP:
            for dirs in dir_names:
                file_path = '{}/{}/{}.txt'.format(self.path_save, dirs, dirs)  # 与生成的txt文件一致

                f = open(file_path)
                lines = f.readlines()
                print("Images Count: {}".format(len(lines)))
                for line in lines:
                    msk_name = '{}/{}/mask/{}.png'.format(self.path_save, dirs, line.strip()[:-4])
                    trimap_name = '{}/{}/trimap/{}.png'.format(self.path_save, dirs, line.strip()[:-4])

                    trimap_name_t = trimap_name.split("/")[:-1]
                    trimap_path = "/".join(trimap_name_t)
                    self.safe_makedirs(trimap_path)

                    msk = cv2.imread(msk_name, 0)
                    trimap = self.erode_dilate(msk, size=(self.KSIZE, self.KSIZE))

                    print("Write to {}".format(trimap_name))
                    cv2.imwrite(trimap_name, trimap)

    def generate_alpha(self):
        """
        由mask图复制得到alpha图
        """
        dir_names = ['train', 'eval', 'test']
        if self.GENERATE_ALPHA:
            for dirs in dir_names:
                path_mask = '{}/{}/mask'.format(self.path_save, dirs)
                path_alpha = '{}/{}/alpha'.format(self.path_save, dirs)
                print('Copying {} mask into alpha...'.format(dirs))
                copy_tree(path_mask, path_alpha)

    def generate_debug(self):
        """
        生成调试用数据集

        选取的数据集：
            train：[alpha, clip_img, mask, matting, trimap]/1803201916/[matting, clip]_00000000/*.png
            eval:  [alpha, clip_img, mask, matting, trimap]/1803232218/[matting, clip]_00000000/*.png
            test:  [alpha, clip_img, mask, matting, trimap]/1803250146/[matting, clip]_00000000/*.png
        """
        dir_names = ['train', 'eval', 'test']
        list_data_type = ['alpha', 'clip_img', 'mask', 'matting', 'trimap']
        dict_dirs = {'train': '1803201916', 'eval': '1803232218', 'test': '1803250146'}
        dict_mt = {'alpha': 'matting', 'clip_img': 'clip', 'mask': 'matting', 'matting': 'matting', 'trimap': 'matting'}
        if self.GENERATE_DEBUG:
            for dirs in dir_names:
                for dt in list_data_type:
                    path_src = '{}/{}/{}/{}/{}_00000000'.format(self.path_save, dirs, dt, dict_dirs[dirs], dict_mt[dt])
                    path_dst = '{}/{}/{}/{}/{}_00000000'.format(self.path_debug, dirs, dt, dict_dirs[dirs], dict_mt[dt])
                    copy_tree(path_src, path_dst)
            self.generate_txt(self.path_debug)


if __name__ == '__main__':
    gd = GenerateData()
    gd.run()
